# coding:utf8
import visdom
import time
import numpy as np
from torchvision import transforms as T
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder


class Visualizer(object):
    """
    封装了visdom的基本操作，但仍然可以通过`self.vis.function`，或者self.function调用原生的visdom接口
    例如
    self.text('hello visdom')
    self.histogram(t.randn(1000))
    self.line(t.arange(0, 10), t.arange(1, 11))
    """

    def __init__(self, env='default'):
        self.vis = visdom.Visdom(env=env)
        self.index = {}
        self.log_text = ''

    def reinit(self, env='default'):
        """
        修改visdom的配置
        """
        self.vis = visdom.Visdom(env=env)
        return self

    def plot_many(self, d):
        """
        一次plot多个
        @params d: dict (name,value) i.e. plot_many({'loss': 1.00, 'acc': 2.0})
        """
        for k, v in d.items():
            self.plot(k, v)

    def img_many(self, d):
        for k, v in d.items():
            self.img(k, v)

    def plot(self, name, y):
        """
        self.plot('loss',1.00)
        :param name: pane和title的名字
        :param y: loss value
        :return:
        """
        x = self.index.get(name, 0)
        if name == 'train-loss':
            x_label = 'batch'
            y_label = 'loss value'

        if name == 'validate-accuracy':
            x_label = 'epoch'
            y_label = 'accuracy'

        self.vis.line(Y=np.array([y]), X=np.array([x]),
                      win=name,
                      opts=dict(title=name, xlabel=x_label, ylabel=y_label),
                      update=None if x == 0 else 'append'
                      )
        self.index[name] = x + 1

    def img(self, name, img_):
        """
        self.img('input_img',t.Tensor(64,64))
        self.img('input_imgs',t.Tensor(3,64,64))
        self.img('input_imgs',t.Tensor(100,1,64,64))
        """
        self.vis.images(img_.cpu().numpy(),
                        win=name,
                        opts=dict(title=name)
                        )

    def log(self, info, win='log_text'):
        """
        self.log({'loss':1,'lr':0.0001})
        """
        self.log_text += ('[{time}] {info} <br>'.format(
                            time=time.strftime('%m-%d %H:%M:%S'),
                            info=info))
        self.vis.text(self.log_text, win)

    def show_confusion_matrix(self, name, cm):
        x = self.index.get(name, 0)

        self.vis.line(
            [[cm[0][0], cm[0][1], cm[1][0], cm[1][1]]],
            [x],
            win=name,
            opts=dict(
                title=name,
                xlabel='epoch',
                ylabel='sample number',
                legend=["TP", "FP", "FN", "TN"],
                dash=np.array(['solid', 'dashdot', 'dashdot', 'solid']),
                linecolor=np.array([[255, 0, 0], [255, 0, 0], [0, 0, 255], [0, 0, 255]])
            ),
            update=None if x == 0 else 'append'
        )
        self.index[name] = x + 1

trans = T.Compose([T.Resize([224, 224]),T.ToTensor()])
imageset = ImageFolder(r'E:\pyprojects\googlenet\data\train', transform=trans)
dataloader = DataLoader(imageset, batch_size=24, shuffle=True)
dataiter = iter(dataloader)
imgs, labels = next(dataiter)
visulizer = Visualizer()
visulizer.img('img', imgs)
